import torch
from layer import GraphConvolution
import torch.nn as nn
import torch.nn.functional as F


class H_GAT(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_classes=2):
        super(H_GAT, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_classes = num_classes
        # 卷积层
        self.conv1 = GraphConvolution(input_dim)
        # 全连接层
        self.fc1 = nn.Linear(291 * hidden_dim, hidden_dim // 2)
        self.fc2 = nn.Linear(hidden_dim // 2, hidden_dim // 4)
        self.fc3 = nn.Linear(hidden_dim // 4, num_classes)

    def forward(self, W, H, n, path_threshold):
        node_H, path_H, net_H = F.relu(self.conv1.forward(W, H, n, path_threshold))
        path_H = path_H[, 0:100, ]
        net_H = net_H[, 0:30, ]
        H_target = torch.cat((node_H, path_H, net_H), dim=1)
        fc1 = F.relu(self.fc1(H_target))
        fc2 = F.relu(self.fc2(fc1))
        logits = self.fc3(fc2)
        return logits
